package com.ztom.agent.trace;

import com.ztom.agent.util.TraceAndSpanIdGenerator;
import lombok.Data;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;

/**
 * @author ZhangTao
 */
@Data
public class Span {
    /**
     * 链路ID
     */
    private String traceId;
    /**
     * 节点ID
     */
    private String spanId;
    /**
     * 父节点ID
     */
    private String parentSpanId;
    /**
     * 节点名称
     */
    private String spanName;
    /**
     * 是否取样
     */
    private boolean sampleable;
    /**
     * 节点类型
     */
    private SpanType spanType;
    /**
     * 节点起始时间毫秒值
     */
    private long spanStartTimeEpochTimeMillis;
    /**
     * 节点起始时间纳秒值
     */
    private long spanStartTimeNanos;
    /**
     * 根节点类型, 1=>根节点, 0=>子节点
     */
    private int rootType = 0;
    /**
     * 消耗时间
     */
    private Integer timeCost;
    /**
     * 未捕获的异常
     */
    private Throwable uncaughtException;
    /**
     * 捕获的异常列表
     */
    private List<Throwable> caughtExceptionList;
    /**
     *
     */
    private boolean rootSpanHasSubSpanError = false;
    /**
     * 扩展信息
     */
    private Map<String, Object> extraDataMap;

    public void addExtraData(String key, Object value) {
        if (extraDataMap == null) {
            extraDataMap = new HashMap<>();
        }
        extraDataMap.put(key, value);
    }

    public enum SpanType {
        /**
         * http请求
         */
        HTTP_REQUEST,

        CLIENT,

        /**
         * 业务类
         */
        LOCAL_ONLY,
        JEDIS_POOL,
        DRUID_POOL,
        DBCP2_POOL,
        TOMCAT_JDBC_POOL,
        PSD_CONNECTION_POOL,
        OK_HTTP,
        COMMONS_HTTP,
        SPRING_HTTP,
        JEDIS_CMD,
        JDBC_EXECUTE,
        DUBBO_CONSUMER,
        DUBBO_PROVIDER,
        MONGO,
        UNKNOWN
    }

    private Span() {
    }

    public Span(String traceId, String parentSpanId, String spanId, String spanName, boolean sampleable,
                SpanType spanType, long spanStartTimeEpochTimeMillis, long spanStartTimeNanos, Integer timeCost) {

        if (traceId == null) {
            throw new IllegalArgumentException("traceId cannot be null");
        }

        if (spanId == null) {
            throw new IllegalArgumentException("spanId cannot be null");
        }

        if (spanName == null) {
            throw new IllegalArgumentException("spanName cannot be null");
        }

        this.traceId = traceId;
        this.spanId = spanId;
        this.parentSpanId = parentSpanId;
        this.spanName = spanName;
        this.sampleable = sampleable;
        this.spanStartTimeEpochTimeMillis = spanStartTimeEpochTimeMillis;
        this.spanStartTimeNanos = spanStartTimeNanos;

        this.timeCost = timeCost;

        if (spanType == null) {
            spanType = SpanType.UNKNOWN;
        }

        this.spanType = spanType;
    }

    /**
     * 给新链路生成根节点
     *
     * @param spanName
     * @param spanType
     * @return
     */
    public static Builder generateRootSpanForNewTrace(String spanName, SpanType spanType) {
        return Span.newBuilder(spanName, spanType);
    }

    /**
     * 生成子节点
     *
     * @param spanName
     * @param spanType
     * @return
     */
    public Span generateChildSpan(String spanName, SpanType spanType) {
        return Span.newBuilder(this)
                .withParentSpanId(this.spanId)
                .withSpanName(spanName)
                .withSpanId(TraceAndSpanIdGenerator.generateId())
                .withSpanStartTimeEpochTimeMillis(System.currentTimeMillis())
                .withSpanStartTimeNanos(System.nanoTime())
                .withTimeCost(null)
                .withSpanPurpose(spanType)
                .build();
    }


    public static Builder newBuilder(String spanName, SpanType spanType) {
        return new Builder(spanName, spanType);
    }


    public static Builder newBuilder(Span copy) {
        Builder builder = new Builder(copy.spanName, copy.spanType);
        builder.traceId = copy.traceId;
        builder.spanId = copy.spanId;
        builder.parentSpanId = copy.parentSpanId;
        builder.sampleable = copy.sampleable;
        builder.spanStartTimeEpochTimeMillis = copy.spanStartTimeEpochTimeMillis;
        builder.spanStartTimeNanos = copy.spanStartTimeEpochTimeMillis;
        builder.timeCost = copy.timeCost;
        return builder;
    }


    public boolean sampleable() {
        return sampleable;
    }


    public void complete() {
        if (this.timeCost != null) {
            throw new IllegalStateException("This Span is already completed!!");
        }

        this.timeCost = (int) (System.currentTimeMillis() - spanStartTimeEpochTimeMillis);
    }


    public boolean hasComplete() {
        return timeCost != null;
    }


    @Override
    public boolean equals(Object o) {
        if (this == o) {
            return true;
        }
        if (!(o instanceof Span)) {
            return false;
        }
        Span span = (Span) o;
        return sampleable == span.sampleable &&
                spanStartTimeEpochTimeMillis == span.spanStartTimeEpochTimeMillis &&
                spanType == span.spanType &&
                Objects.equals(traceId, span.traceId) &&
                Objects.equals(spanId, span.spanId) &&
                Objects.equals(parentSpanId, span.parentSpanId) &&
                Objects.equals(spanName, span.spanName) &&
                Objects.equals(timeCost, span.timeCost);
    }

    @Override
    public int hashCode() {
        return Objects.hash(traceId, spanId, parentSpanId, spanName, sampleable, spanType, spanStartTimeEpochTimeMillis, spanStartTimeNanos, timeCost);
    }

    public static final class Builder {

        private String traceId;
        private String spanId;
        private String parentSpanId;
        private String spanName;
        private boolean sampleable = true;
        private Long spanStartTimeEpochTimeMillis;
        private Long spanStartTimeNanos;
        private Integer timeCost;
        private SpanType spanType;

        private Builder(String spanName, SpanType spanType) {
            this.spanName = spanName;
            this.spanType = spanType;
        }

        public Builder withTraceId(String traceId) {
            this.traceId = traceId;
            return this;
        }

        public Builder withSpanId(String spanId) {
            this.spanId = spanId;
            return this;
        }

        public Builder withParentSpanId(String parentSpanId) {
            this.parentSpanId = parentSpanId;
            return this;
        }

        public Builder withSpanName(String spanName) {
            this.spanName = spanName;
            return this;
        }

        public Builder withSampleable(boolean sampleable) {
            this.sampleable = sampleable;
            return this;
        }

        public Builder withSpanStartTimeEpochTimeMillis(Long spanStartTimeEpochMicros) {
            this.spanStartTimeEpochTimeMillis = spanStartTimeEpochMicros;
            return this;
        }

        public Builder withSpanStartTimeNanos(Long spanStartTimeNanos) {
            this.spanStartTimeNanos = spanStartTimeNanos;
            return this;
        }

        public Builder withTimeCost(Integer durationNanos) {
            this.timeCost = durationNanos;
            return this;
        }

        public Builder withSpanPurpose(SpanType spanType) {
            this.spanType = spanType;
            return this;
        }

        public Span build() {
            if (traceId == null) {
                traceId = TraceAndSpanIdGenerator.generateId();
            }

            if (spanId == null) {
                spanId = TraceAndSpanIdGenerator.generateId();
            }

            if (spanStartTimeEpochTimeMillis == null) {
                spanStartTimeEpochTimeMillis = System.currentTimeMillis();
            }
            if (spanStartTimeNanos == null) {
                spanStartTimeNanos = System.nanoTime();
            }

            return new Span(traceId, parentSpanId, spanId, spanName, sampleable, spanType, spanStartTimeEpochTimeMillis, spanStartTimeNanos, timeCost);
        }
    }
}
